#!/usr/bin/env python3

import os
import pandas as pd
import glob
import subprocess

script_dir = os.path.dirname(os.path.abspath(__file__))

# List of files
files = glob.glob(f"{script_dir}/output/**/*.ttir", recursive=True)

# Initialize a list to hold the mappings
mappings = []


def extract_gpu(s):
    gpus = {"blackwell", "hopper", "ampere"}
    for gpu in gpus:
        if gpu not in s:
            continue
        sub_s = s.split(gpu)
        example = sub_s[0][:-1]
        return True, example, gpu

    print(f"Warning: {s} not handled")
    return False, "", ""


# Process each file to extract example and kernel
for file_name in files:
    path_parts = file_name.strip().split('output/')
    path_parts = path_parts[1].split('/')
    success, example, gpu = extract_gpu(path_parts[0])
    if success:
        ttir = os.path.abspath(file_name)
        ttir_dir = os.path.dirname(ttir)
        kernel = os.path.splitext(os.path.basename(ttir))[0]

        mlir_file_path = os.path.join(os.path.dirname(file_name), "dump.mlir")
        if not os.path.exists(mlir_file_path):
            continue
        extract_passes_script = os.path.join(script_dir, 'extract_passes.py')
        result = subprocess.run(
            ['python3', extract_passes_script, mlir_file_path])
        assert result.returncode == 0, f"{extract_passes_script} failed"

        info_dir = os.path.join(ttir_dir, kernel)
        assert os.path.exists(info_dir), f"{info_dir} doesn't exist"
        passes = open(os.path.join(info_dir, "passes.txt"), 'r').read().strip()
        passes_changed = open(os.path.join(
            info_dir, "passes_changed.txt"), 'r').read().strip()

        mappings.append((example, gpu, kernel, passes, passes_changed))

# Create a DataFrame from the mappings
df = pd.DataFrame(mappings, columns=[
                  'Example', 'GPU', 'Kernel', 'Passes', 'Passes Changed'])

# Print the DataFrame to standard output
print(df)

# Optionally, save to CSV if needed
output_csv = 'kernel_passes.csv'
df.to_csv(output_csv, index=False)
print(f'Successfully written mappings to {output_csv}')
